import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
from openpyxl import Workbook

# 1. 读取训练数据
train_data = pd.read_excel('合并数据集.xlsx')

# 提取所需特征和目标变量
material = pd.Categorical(train_data['材料']).codes  # 材料：材料1、材料2、材料3、材料4
temperature = train_data['温度，oC']                  # 温度：25、50、70、90摄氏度
frequency = train_data['频率，Hz']                    # 频率：50000—500000 Hz
core_loss = train_data['磁芯损耗，w/m3']             # 磁芯损耗，目标变量
waveform = pd.Categorical(train_data['励磁波形']).codes  # 励磁波形：正弦波、三角波、梯形波

# 提取磁通密度峰值（假设磁通密度从第6列到最后列）
magnetic_density = train_data.iloc[:, 5:].values      # 提取磁通密度数据
Bm_train = np.max(magnetic_density, axis=1)           # 提取每行（每个样本）的磁通密度峰值

# 构建特征矩阵（输入变量），加入磁通密度峰值
X_train = np.column_stack([material, temperature, frequency, waveform, Bm_train])

# 2. 构建线性回归模型来预测磁芯损耗
model = LinearRegression()

# 训练模型
model.fit(X_train, core_loss)

# 输出模型的评分
score = model.score(X_train, core_loss)
print(f'线性回归模型的R²得分: {score}')

# 3. 读取测试数据
test_data = pd.read_excel('附件三（测试集）.xlsx')

# 提取测试数据中的特征
material_test = pd.Categorical(test_data['磁芯材料']).codes   # 材料：材料1、材料2、材料3、材料4
temperature_test = test_data['温度，oC']                      # 温度
frequency_test = test_data['频率，Hz']                        # 频率
waveform_test = pd.Categorical(test_data['励磁波形']).codes   # 励磁波形

# 提取测试数据中的磁通密度峰值（假设磁通密度从第6列到最后列）
magnetic_density_test = test_data.iloc[:, 5:].values      # 提取磁通密度数据
Bm_test = np.max(magnetic_density_test, axis=1)           # 提取每行（每个样本）的磁通密度峰值

# 构建测试数据的特征矩阵，加入磁通密度峰值
X_test = np.column_stack([material_test, temperature_test, frequency_test, waveform_test, Bm_test])

# 4. 使用训练好的线性回归模型对附件三中的测试数据进行预测
core_loss_pred = model.predict(X_test)

# 5. 保存预测结果为 Excel 文件
predicted_core_loss_table = pd.DataFrame(core_loss_pred, columns=['PredictedCoreLoss'])

# 保存到 Excel 文件
predicted_core_loss_table.to_excel('问题四预测磁芯损耗结果包含磁通密度线性回归版做法.xlsx', index=False)
print('预测结果已成功保存为 问题四预测磁芯损耗结果包含磁通密度线性回归版做法.xlsx')

# 6. 绘制预测结果
plt.figure(figsize=(8, 6))
plt.plot(core_loss_pred, 'r-', lw=1.5)
plt.xlabel('样本编号')
plt.ylabel('预测磁芯损耗 (W/m^3)')
plt.title('磁芯损耗预测结果 - 线性回归')
plt.grid(True)
plt.legend(['预测磁芯损耗'])
plt.show()
